from fastapi import HTTPException
from starlette.responses import JSONResponse
from models.corpus import Corpus

async def query_all_corpus(question: str = None, answer: str = None, status: bool = None, from_model: str = None, from_class: str = None) -> list[Corpus]:
    query = Corpus.filter()

    if question:
        query = query.filter(question=question)
    
    if answer:
        query = query.filter(answer=answer)

    if status is not None:
        query = query.filter(status=status)
    
    if from_model:
        query = query.filter(from_model=from_model)

    if from_class:
        query = query.filter(from_class=from_class)

    corpus = await query
    if not corpus:
        raise HTTPException(status_code=403, detail="未找到匹配语料")
    return corpus

async def add_corpus(question: str, answer: str, status: bool, from_model: str, from_class: str) -> Corpus:
    corpus = await Corpus.create(
        question=question,
        answer=answer,
        status=status,
        from_model=from_model,
        from_class=from_class
    )
    if corpus:
        return {"msg": "创建语料成功"}
    else:
        return {"msg": "创建语料失败"}

async def delete_corpus_by_id(corpus_id: int) -> Corpus:
    corpus = await Corpus.get_or_none(id=corpus_id)
    if corpus:
        await corpus.delete()
        return {"msg": "删除语料成功"}
    else:
        return {"msg": "删除语料失败"}

async def update_corpus_by_id(id: int, question: str, answer: str, status: bool, from_model: str, from_class: str) -> Corpus:
    corpus = await Corpus.get_or_none(id=id)
    if corpus:
        await corpus.update_from_dict(
            {
                "question": question,
                "answer": answer,
                "status": status,
                "from_model": from_model,
                "from_class": from_class
            }
        )
        await corpus.save()
        return {"msg": "更新语料成功"}
    else:
        return {"msg": "更新语料失败"}